'''
	Copyright (C) 2010  Pier Maria Mattioli
	
	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.
	
	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.
	
	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
	
'''

import os
import time
import hash_ring
import threading
import redis
import marshal
import fnmatch
import random
import multiprocessing
import Queue

def node_exp_handler(default):
	def decorator(func):
		def wrapper(*pargs, **kwargs):
			try:
				return func(*pargs, **kwargs)
			except Exception, err:	
				return default[0]
		return wrapper
	return decorator


class RedisCli(redis.Redis):
	def __init__(self, connection):
		self.encoding = 'utf8'
		self.errors = 'strict'
		self.connection = connection
		self.subscribed = False

class Nodes():

	def __init__(self):
		self.pool = redis.ConnectionPool()
	
	@node_exp_handler((None,))
	def get_node(self, node_id):
		host, port, db = node_id.split(':')
		connection = self.pool.get_connection(host, int(port), int(db), None, None)
		return RedisCli(connection)
		
	@node_exp_handler((None,))
	def do_set(self, node, key, value):
		value = marshal.dumps(value)
		return node.set(key, value)

	@node_exp_handler((None,))
	def do_setnx(self, node, key, value):
		value = marshal.dumps(value)
		resp = node.setnx(key, value)
		return resp

	@node_exp_handler((None,))
	def do_get(self, node, key, default=None):
		value = node.get(key)
		if value:
			return marshal.loads(value)
		else:
			return default
	
	@node_exp_handler((None,))
	def do_delete(self, node, key, empty=None):
		return node.delete(key)

	@node_exp_handler(([],))
	def do_keys(self, node, pattern):
		return node.keys(pattern)

	@node_exp_handler(({},))
	def do_mget(self, node, mapping):
		keys = [m[0] for m in mapping]
		values = node.mget(keys)
		out = {}
		for pos,v in enumerate(values):
			k = keys[pos]
			if v:
				out[k] = marshal.loads(v)
		return out

	@node_exp_handler((None,))			
	def do_mset(self, node, mapping):
		redis_mapping = {}
		for k, v in mapping:
			redis_mapping[k] = marshal.dumps(v)
		return node.mset(redis_mapping)

	@node_exp_handler((None,))			
	def do_mdelete(self, node, mapping):	
		for ele in mapping:
			node.delete(ele[0])


class Threads():
	def __init__(self):
		self.threads = []
		self.halt = False
	
	def stop_threads(self):
		self.halt = True
		for th, w_queue, r_queue in self.threads:
			th.join()
	
	def reap_thread(self):
		for pos, (th, w_queue, r_queue) in enumerate(self.threads):
			if not th.isAlive():
				del self.threads[pos]
	
	def start_thread(self):
		w_queue = Queue.Queue()
		r_queue = Queue.Queue()
		th = threading.Thread(target=self.worker, args=(w_queue,r_queue))
		self.threads.append((th,w_queue,r_queue))
		th.daeomn = True
		th.start()
		
	def map(self, function, datalist):
		self.reap_thread()
		for i in range(len(datalist) - (len(self.threads))):
			self.start_thread()
		
		for pos,data in enumerate(datalist):
			self.threads[pos][1].put((function, data))

		result = []
		for pos,data in enumerate(datalist):
			result.append(self.threads[pos][2].get(block=True, timeout=None))
		return result
			
	def worker(self, r_queue, w_queue):
		while not self.halt:
			try:
				function, data = r_queue.get(block=True, timeout=2)
			except Queue.Empty:
				continue
			result = function(*data)
			w_queue.put(result)

class Cluster():

	def __init__(self, nodes='localhost:6379:0', redundancy=2):
		self.nodes = Nodes()
		self.threads = Threads()
		self.ctrl_key = 'ctr::nodes::all'
		self.redundancy = redundancy
		if isinstance(nodes, str): nodes = [nodes]
		self.ring, self.container = self._init_container(nodes)
		self._check_container()

	def get(self, key, default=None):
		self._check_container()
		nodes = self._get_nodes(key)
		for node in nodes:
			result = self.nodes.do_get(node, key, None)
			if result is not None:
				break
		else:
			result = default
		return result
		
	def _get_ctrl_key(self):
		key = self.ctrl_key
		nodes = self.container.values()
		random.shuffle(nodes)
		for node in nodes:
			result = self.nodes.do_get(node, key, None)
			if result is not None:
				break
		else:
			result = None
		return result
	
	def set(self, key, value):
		self._check_container()
		map_list = self._get_map_list(key, value)
		self.threads.map(self.nodes.do_set, map_list)

	def setnx(self, key, value):
		self._check_container()
		self.get_lock(key)
		result = self._setnx(key, value)
		self.release_lock(key)
		return result
			
	def delete(self, key):
		self._check_container()
		self.get_lock(key)
		map_list = self._get_map_list(key)
		self.threads.map(self.nodes.do_delete, map_list)
		self.release_lock(key)
		
	def keys(self, pattern):
		self._check_container()
		catch_list = [(node, pattern) for node in self.container.values()]
		keys = self.threads.map(self.nodes.do_keys, catch_list)
		[keys[0].extend(k) for k in keys[1:]]
		return list(set(keys[0]))

	def mget(self, mapping):
		self._check_container()
		sharded_mapping = self._get_multi_map_list(mapping)
		result = self.threads.map(self.nodes.do_mget, sharded_mapping)
		out = {}
		key_list = []
		for ele in mapping:
			if isinstance(ele, tuple):
				key, default = ele
			else:
				key = ele
				default = None
			key_list.append(key)
			for rec in result:
				if rec.get(key) is not None:
					out[key] = rec[key]
					break
			else:
				out[key] = default
		return [out[k] for k in key_list]
	
	def mset(self, mapping):
		self._check_container()
		sharded_mapping = self._get_multi_map_list(mapping)
		self.threads.map(self.nodes.do_mset, sharded_mapping)

	def mdelete(self, mapping):
		self._check_container()
		sharded_mapping = self._get_multi_map_list(mapping)
		self.threads.map(self.nodes.do_mdelete, sharded_mapping)

	def catch_lock(self, lock):
		lock = 'lock.%s' % lock
		self._check_container()
		mapping = self._get_map_list(lock, time.time())
		result = self.threads.map(self.nodes.do_setnx, mapping)
		if False in result:
			to_del = [map[:2] for pos, map in enumerate(mapping) if result[pos] == True]
			self.threads.map(self.nodes.do_delete, to_del)
			return False
		return True

	def get_lock(self, lock, times=0, delay=0.00005):
		if times == 0:
			while not self.catch_lock(lock):
				time.sleep(delay)
			return True
				
		else:
			for i in range(times):
				if self.catch_lock(lock):
					return True
				time.sleep(delay)
			return False
			
	def release_lock(self, lock):
		lock = 'lock.%s' % lock
		self._check_container()
		mapping = self._get_map_list(lock, time.time(), overflow=self.redundancy)
		result = self.threads.map(self.nodes.do_delete, mapping)

	def _setnx(self, key, value):
		map_list = self._get_map_list(key, value)
		result = self.threads.map(self.nodes.do_setnx, map_list)
		return result
			
	def _get_multi_map_list(self, mapping):	
		d = {}
		for map in mapping:
			if isinstance(map, tuple):
				node_maps = self._get_map_list(*map)
			else:
				node_maps = self._get_map_list(map)
			for node_map in node_maps:
				node_id = node_map[0]
				if not node_id in d.keys():
					d[node_id] = []
				d[node_id].append(node_map[1:])			
		return d.items()
	
	def _init_container(self, container_ids):
		ring = hash_ring.HashRing(container_ids)
		container = {}
		for node_id in container_ids:
			node = self.nodes.get_node(node_id)
			if node:
				container[node_id] = node
		return ring, container

	def _get_nodes_iterator(self, key):
		iter_ring = self.ring.iterate_nodes(key)
		for node_id in iter_ring:
			yield self.container.get(node_id, None)
			
	def _get_nodes(self, key):
		iter_nodes = self._get_nodes_iterator(key)
		return [iter_nodes.next() for i in range(self.redundancy)]
		
	def _check_container(self):
		cur_container_ids = self._get_ctrl_key()
		old_container_ids = [self.container.keys(), self.redundancy]
		if cur_container_ids != old_container_ids:
			if cur_container_ids is None:
				for node in self.container.values():
					self.nodes.do_set(node, self.ctrl_key, old_container_ids)
			else:
				self.ring, self.container = self._init_container(cur_container_ids[0])
				self.redundancy = int(cur_container_ids[1])

	def _get_map_list(self, key, value=None, overflow=0):
		iter_nodes = self._get_nodes_iterator(key)
		N = min(len(self.container.keys()), (self.redundancy + overflow))
		node_mapping = tuple((iter_nodes.next(), key, value) for i in range(N))
		return node_mapping
	
	def shutdown(self):
		# monkey monkey monkey
		self.threads.stop_threads()
		
class Manager(Cluster):

	def add_node(self, node_id):
		self._get_ctrl_lock()
		container_ids = self._get_ctrl_key()
		new_node = self.nodes.get_node(node_id)
		if container_ids is None:
			result = 'Custer corrupted, good luck !'
		elif node_id not in container_ids[0] and new_node:
			self.redundancy = max(self.redundancy, 2)
			self.init_node(new_node)
			container_ids[0].append(node_id)
			container_ids[1] = self.redundancy
			self.nodes.do_set(new_node, self.ctrl_key, container_ids)
			self._publicate_container(container_ids)
			result = 'done'
		else:
			result = 'nothing to be done'
		self._release_ctrl_lock()
		return result
		
	def del_node(self, node_id):
		self._get_ctrl_lock()
		container_ids = self._get_ctrl_key()
		if self.redundancy == 1:
			result = 'Nooooooooooooooooo'
		elif container_ids is None:
			result = 'Custer corrupted, good luck !'
		elif node_id in container_ids[0]:
			del container_ids[0][container_ids[0].index(node_id)]
			container_ids[1] = min(self.redundancy, len(container_ids[0]))
			self._publicate_container(container_ids)
			result = 'done'
		else:
			result = 'nothing to be done'
		self._release_ctrl_lock()
		return result

	def rebalance(self, node_id=None):
		self._get_ctrl_lock()
		time.sleep(2)
		self._rebalance_container(node_id)
		self._release_ctrl_lock()
		return 'done'

	def change_redundancy(self, new_value):
		self._get_ctrl_lock()
		self.redundancy = min(new_value, len(self.container.keys()))
		self._publicate_container([self.container.keys(), self.redundancy])
		self._release_ctrl_lock()
		return 'done'

	def init_node(self, node):
		keyspace = self.nodes.do_keys(node, '*')
		for k in keyspace:
			self.nodes.do_delete(node, k)

	def _publicate_container(self, container_ids):
		mapping = [(node, self.ctrl_key, container_ids) for node in self.container.values()]
		self.threads.map(self.nodes.do_set, mapping)		

	def _rebalance_container(self, node_id):
		if node_id is None or node_id in self.container.keys():
			ring = self.ring
		else:
			ring = hash_ring.HashRing(self.container.keys() + [node_id])
		keyspace = self.keys('*')
		for key in keyspace:
			if key != self.ctrl_key:
				iter_nodes = ring.iterate_nodes(key)
				node_ids = [iter_nodes.next() for i in range(self.redundancy)]
				if node_id is None\
				or node_id in node_ids:

					self.get_lock(key)
					value = self.get(key, None)
					if value is not None:
						self._setnx(key, value)
					self.release_lock(key)

					mapping = []
					for pos, node in enumerate(self._get_nodes_iterator(key)):
						if pos >= (self.redundancy * 2):
							break
						if pos >= self.redundancy:
							mapping.append((node, key))
					self.threads.map(self.nodes.do_delete, mapping)	
					
	def _get_ctrl_lock(self):
		self.get_lock(self.ctrl_key)
			
	def _release_ctrl_lock(self):
		self.release_lock(self.ctrl_key)